#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt

# Set up Chinese font support
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def mexican_hat(t):
    """Mexican hat wavelet function"""
    return (1 - t**2) * np.exp(-t**2 / 2)

def create_signal(t):
    """Create test signal with two pulses"""
    signal = np.zeros_like(t)
    signal += 2.0 * np.exp(-50 * (t - 2.0)**2)
    signal += 1.5 * np.exp(-10 * (t - 5.0)**2)
    return signal

def run_convolution_demo():
    # Time array
    t = np.linspace(0, 8, 100)
    dt = t[1] - t[0]
    
    # Create signal and wavelet
    signal = create_signal(t)
    t_wavelet = np.linspace(-2, 2, 50)
    wavelet = mexican_hat(t_wavelet)
    
    # Calculate convolution using numpy
    conv_result = np.convolve(signal, wavelet, mode='same') * dt
    
    # Create visualization of signal and wavelet
    plt.figure(figsize=(10, 6))
    plt.subplot(2, 1, 1)
    plt.plot(t, signal, 'b-')
    plt.title('Original Signal')
    plt.xlabel('Time')
    plt.ylabel('Amplitude')
    plt.grid(True)
    
    plt.subplot(2, 1, 2)
    plt.plot(t_wavelet, wavelet, 'r-')
    plt.title('Mexican Hat Wavelet')
    plt.xlabel('Time')
    plt.ylabel('Amplitude')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('signal_wavelet.png')
    plt.close()
    
    # Visualize convolution result
    plt.figure(figsize=(10, 6))
    plt.plot(t, conv_result, 'g-')
    plt.title('Convolution Result')
    plt.xlabel('Time')
    plt.ylabel('Convolution Value')
    plt.grid(True)
    plt.savefig('convolution_result.png')
    plt.close()
    
    # Demonstrate convolution process step by step
    reversed_wavelet = np.flip(wavelet)
    n = len(signal)
    m = len(wavelet)
    
    # Select key positions to show
    positions = [20, 40, 60, 80]
    
    plt.figure(figsize=(10, 12))
    for i, pos in enumerate(positions):
        plt.subplot(len(positions), 1, i+1)
        
        # Plot signal
        plt.plot(t, signal, 'b-', label='Signal')
        
        # Create wavelet at current position
        wavelet_at_pos = np.zeros_like(signal)
        start = max(0, pos - m + 1)
        end = min(pos + 1, m)
        if end > start:
            wavelet_at_pos[start:pos+1] = reversed_wavelet[m-end:m-start]
        
        # Plot wavelet at current position
        plt.plot(t, wavelet_at_pos, 'r-', label='Wavelet')
        
        # Calculate and show product
        product = signal * wavelet_at_pos
        plt.bar(t, product, width=0.08, alpha=0.5, color='g', label='Product')
        
        # Calculate convolution value
        conv_value = np.sum(product) * dt
        plt.title('Step %d: t=%.2f, Conv=%.4f' % (pos, t[pos], conv_value))
        plt.grid(True)
        plt.legend()
    
    plt.tight_layout()
    plt.savefig('convolution_steps.png')
    plt.close()
    
    # Show scaling effects
    scales = [0.5, 1.0, 2.0]
    plt.figure(figsize=(10, 12))
    
    for i, scale in enumerate(scales):
        t_scaled = np.linspace(-3*scale, 3*scale, 100)
        scaled_wavelet = mexican_hat(t_scaled/scale) / scale
        scaled_conv = np.convolve(signal, scaled_wavelet, mode='same') * dt
        
        plt.subplot(len(scales), 1, i+1)
        plt.plot(t, scaled_conv, 'b-')
        plt.title('Scale = ' + str(scale))
        plt.xlabel('Time')
        plt.ylabel('Convolution Value')
        plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('wavelet_scales.png')
    plt.close()
    
    print("\nGenerated visualization files:")
    print("1. signal_wavelet.png - Original signal and wavelet")
    print("2. convolution_result.png - Final convolution result")
    print("3. convolution_steps.png - Step-by-step convolution process")
    print("4. wavelet_scales.png - Convolution with different scales")

def explain_convolution():
    print("\nCONVOLUTION CALCULATION PROCESS:")
    print("===============================")
    print("\n1. BASIC PRINCIPLE:")
    print("   Convolution combines two functions to produce a third function")
    print("   that represents how the shape of one is modified by the other.")
    print("\n2. STEP-BY-STEP CALCULATION:")
    print("   a. Reverse the wavelet function")
    print("   b. Slide the reversed wavelet across the input signal")
    print("   c. At each position, multiply corresponding signal and wavelet values")
    print("   d. Sum all products to get the convolution value at that position")
    print("\n3. MEXICAN HAT WAVELET:")
    print("   - Formula: (1 - t^2) * exp(-t^2 / 2)")
    print("   - Second derivative of Gaussian function")
    print("   - Good for detecting edges and changes in signals")
    print("\n4. RESULT INTERPRETATION:")
    print("   - Positive peaks: Similarity between signal and wavelet")
    print("   - Negative peaks: Inverse similarity")
    print("   - Narrow pulses create paired positive-negative responses")
    print("   - Wide pulses create broader responses")
    print("\n5. SCALING EFFECTS:")
    print("   - Small scale (0.5): Detects fine details and rapid changes")
    print("   - Medium scale (1.0): Balanced detection")
    print("   - Large scale (2.0): Captures broader features")

def main():
    print("Running Mexican hat wavelet convolution demonstration...")
    run_convolution_demo()
    explain_convolution()
    print("\nDemonstration completed!")

if __name__ == "__main__":
    main()